import argparse
import os
import re
from transforms3d.quaternions import quat2mat
import numpy as np

def prepare_azure_timestamp(input, output):
    print(input)
    with open(input, 'r') as f:
        data = f.read()
    timestamps_str = data.split(',')
    timestamps = [int(x.replace('.','').ljust(19,'0')) for x in timestamps_str]
    with open(output, 'w') as f:
        f.write(','.join([str(x) for x in timestamps]))
    return timestamps
    

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--date', type=str)
    args = parser.parse_args()
    date = args.date
    if date is None:
        raise ValueError('Please input date')
    vts_path = os.path.join('/share/hhodataset/VTS', date)
    
    for folder_name in os.listdir(vts_path):
        try:
            if not re.match(r"\d{3}$", folder_name):
                continue
            vts_data_path = os.path.join(vts_path, folder_name, 'VTS_data.npz')
            azure2_path = os.path.join('/share/hhodataset/azure2', date, folder_name)
            azure3_path = os.path.join('/share/hhodataset/azure3', date, folder_name)
            azure4_path = os.path.join('/share/hhodataset/azure4', date, folder_name)
            azure5_path = os.path.join('/share/hhodataset/azure5', date, folder_name)
            azure_paths = ["0","1",azure2_path, azure3_path, azure4_path, azure5_path]
            if not os.path.exists(vts_data_path):
                print('{} :VTS data not exist '.format(vts_data_path))
                continue
            if "aligned_objposes.npy" in os.listdir(os.path.join(vts_path, folder_name)):
                print('{} :aligned_objposes.npy exist '.format(vts_data_path))
                continue
            if "azure2.mkv" not in os.listdir(azure2_path) and "azure2.txt" not in os.listdir(azure2_path):
                print('{} :azure2.mkv not exist '.format(azure2_path))
                continue
            if "azure3.mkv" not in os.listdir(azure3_path) and "azure3.txt" not in os.listdir(azure3_path):
                print('{} :azure3.mkv not exist '.format(azure3_path))
                continue
            if "azure4.mkv" not in os.listdir(azure4_path) and "azure4.txt" not in os.listdir(azure4_path):
                print('{} :azure4.mkv not exist '.format(azure4_path))
                continue
            if "azure5.mkv" not in os.listdir(azure5_path) and "azure5.txt" not in os.listdir(azure5_path):
                print('{} :azure5.mkv not exist '.format(azure5_path))
                continue
            print('start process {}'.format(folder_name))
            
            timestamps_2 = prepare_azure_timestamp(os.path.join(azure2_path, 'azure2.txt'), os.path.join(azure2_path, 'azure2_timestamp.txt'))
            timestamps_3 = prepare_azure_timestamp(os.path.join(azure3_path, 'azure3.txt'), os.path.join(azure3_path, 'azure3_timestamp.txt'))
            timestamps_4 = prepare_azure_timestamp(os.path.join(azure4_path, 'azure4.txt'), os.path.join(azure4_path, 'azure4_timestamp.txt'))
            timestamps_5 = prepare_azure_timestamp(os.path.join(azure5_path, 'azure5.txt'), os.path.join(azure5_path, 'azure5_timestamp.txt'))
            
            data = np.load(vts_data_path, allow_pickle=True)
            VTS_data = data["data"].item()
            timestamps_p1 = VTS_data["person1_timestamp"]
            timestamps_p2 = VTS_data["person2_timestamp"]
            timestamps_obj = VTS_data["rigid_timestamp"]
            
            labels = VTS_data["/labels"]
            N_obj = {}
            for names in labels:
                for name in names:
                    if not name in N_obj:
                        N_obj[name] = 0
                    N_obj[name] += 1
            N_obj["action1"] = 0
            max_n = 0
            obj_name = ""
            for name, N in N_obj.items():
                if N > max_n:
                    max_n = N
                    obj_name = name
            
            threshold =40000000
                       
            # 1696250851 792373400
            # 1696250851 768523700
            # 1696250851 328248300
            # 1696250851 790335000
            # 1696250851 797421900
            # 1696250851 787918700
            # 1696250851 788244600
            
            # 1696250838 852216500
            # 1696250836 601587000
            
            # 1696250851 865238700
            # 1696250851 848567500
            # 1696250851 328248300
            # 1696250851 846571400
            # 1696250851 863720600
            # 1696250851 864945400
            # 1696250851 865257500
            
            
            # 1699533646 000
            #             40000000
            # 1696250840.1486802
            N_2 = len(timestamps_2)
            N_3 = len(timestamps_3)
            N_4 = len(timestamps_4)
            N_5 = len(timestamps_5)
            N_obj = len(timestamps_obj)
            N_p1 = len(timestamps_p1)
            N_p2 = len(timestamps_p2)
            
            p_3, p_4, p_5, p_obj, p_p1, p_p2 = 0, 0, 0, 0, 0, 0
            print(N_2, N_3, N_4, N_5, N_obj, N_p1, N_p2)
            idxs = []
            obj2world_list = []
            
            for p_2 in range(N_2):
                t_2 = timestamps_2[p_2]
                while((p_3+1 < N_3) and (abs(t_2 - timestamps_3[p_3+1]) < abs(t_2 - timestamps_3[p_3]))):
                    p_3 += 1
                while((p_4+1 < N_4) and (abs(t_2 - timestamps_4[p_4+1]) < abs(t_2 - timestamps_4[p_4]))):
                    p_4 += 1
                while((p_5+1 < N_5) and (abs(t_2 - timestamps_5[p_5+1]) < abs(t_2 - timestamps_5[p_5]))):
                    p_5 += 1
                while((p_obj+1 < N_obj) and (abs(t_2 - timestamps_obj[p_obj+1]) < abs(t_2 - timestamps_obj[p_obj]))):
                    p_obj += 1
                while((p_p1+1 < N_p1) and (abs(t_2 - timestamps_p1[p_p1+1]) < abs(t_2 - timestamps_p1[p_p1]))):
                    p_p1 += 1
                while((p_p2+1 < N_p2) and (abs(t_2 - timestamps_p2[p_p2+1]) < abs(t_2 - timestamps_p2[p_p2]))):
                    p_p2 += 1
                
                # print(t_2)
                # print(timestamps_3[p_3])
                # print(timestamps_4[p_4])
                # print(timestamps_5[p_5])
                # print(timestamps_obj[p_obj])
                # print(timestamps_p1[p_p1])
                # print(timestamps_p2[p_p2])
                flag = True
                if abs(t_2 - timestamps_3[p_3]) > threshold:
                    flag = False
                if abs(t_2 - timestamps_4[p_4]) > threshold:
                    flag = False
                if abs(t_2 - timestamps_5[p_5]) > threshold:
                    flag = False
                if abs(t_2 - timestamps_obj[p_obj]) > threshold:
                    flag = False
                if abs(t_2 - timestamps_p1[p_p1]) > threshold:
                    flag = False
                if abs(t_2 - timestamps_p2[p_p2]) > threshold:
                    flag = False
                
                if not flag:
                    print("[error in pairing data] at {} : {}".format(folder_name, p_2))
                    continue
                rigid_poses = VTS_data["/rigid"][p_obj]
                devices = labels[p_obj]
                obj2world = None
                for i, name in enumerate(devices):
                    if name == obj_name:
                        rigid_pose = rigid_poses[i]
                        obj2world = np.eye(4)
                        obj2world[:3,3] = rigid_pose["position"]
                        obj2world[:3,:3] = quat2mat(rigid_pose["orientation"])
                        break
                if obj2world is None:
                    continue
                idxs.append(",".join([str(x) for x in [p_2, p_3, p_4, p_5, p_p1, p_p2]]))
                obj2world_list.append(obj2world)
            with open(os.path.join(vts_path, folder_name, "aligned_frame_ids.txt"), 'w') as f:
                f.write("\n".join(idxs))
            # save obj2world as npy aligned_objposes.npy
            np.save(os.path.join(vts_path, folder_name, "aligned_objposes.npy"), np.array(obj2world_list))
            
            # exit()
            
        except Exception as e:
            print("Wrong at {} : {}".format(folder_name, str(e)))
            continue